#!/usr/bin/env python3
"""
提取简单恶意包样本的脚本
从恶意包中筛选出只包含1个或2个JS文件的简单样本
"""

import os
import shutil
from pathlib import Path
from collections import defaultdict
import json

class SimplePackageExtractor:
    def __init__(self, source_dir="E:/extracted_packages/malicious", 
                 output_dir="E:/simplePackages"):
        """
        初始化提取器
        
        Args:
            source_dir: 恶意包源目录
            output_dir: 输出目录
        """
        self.source_dir = Path(source_dir)
        self.output_dir = Path(output_dir)
        
        # 创建输出目录结构
        self.single_js_dir = self.output_dir / "single_js_packages"
        self.double_js_dir = self.output_dir / "double_js_packages"
        
        # 创建目录
        self.single_js_dir.mkdir(parents=True, exist_ok=True)
        self.double_js_dir.mkdir(parents=True, exist_ok=True)
        
        # 统计信息
        self.stats = {
            "total_packages": 0,
            "single_js_packages": 0,
            "double_js_packages": 0,
            "skipped_packages": 0,
            "error_packages": 0
        }
        
        # 详细记录
        self.extraction_log = []

    def count_js_files(self, package_dir):
        """
        递归统计包中的JS文件数量
        
        Args:
            package_dir: 包目录路径
            
        Returns:
            tuple: (js文件数量, js文件路径列表)
        """
        js_files = []
        
        try:
            for root, dirs, files in os.walk(package_dir):
                for file in files:
                    if file.endswith('.js'):
                        js_files.append(Path(root) / file)
            
            return len(js_files), js_files
            
        except Exception as e:
            print(f"❌ 统计JS文件失败 {package_dir}: {e}")
            return 0, []

    def copy_package(self, source_package, target_dir, category):
        """
        复制包到目标目录
        
        Args:
            source_package: 源包路径
            target_dir: 目标目录
            category: 分类 ('single' 或 'double')
        """
        try:
            package_name = source_package.name
            target_path = target_dir / package_name
            
            # 如果目标已存在，添加数字后缀
            counter = 1
            while target_path.exists():
                target_path = target_dir / f"{package_name}_{counter}"
                counter += 1
            
            # 复制整个包目录
            shutil.copytree(source_package, target_path)
            
            print(f"✅ 复制成功: {package_name} -> {category}_js_packages")
            return True
            
        except Exception as e:
            print(f"❌ 复制失败 {source_package.name}: {e}")
            return False

    def extract_simple_packages(self):
        """
        提取简单包的主函数
        """
        print("🚀 开始提取简单恶意包样本...")
        print(f"源目录: {self.source_dir}")
        print(f"输出目录: {self.output_dir}")
        print("-" * 60)
        
        if not self.source_dir.exists():
            print(f"❌ 源目录不存在: {self.source_dir}")
            return
        
        # 获取所有子目录（恶意包）
        malicious_packages = [d for d in self.source_dir.iterdir() if d.is_dir()]
        
        print(f"📦 发现 {len(malicious_packages)} 个恶意包")
        
        for package_dir in malicious_packages:
            self.stats["total_packages"] += 1
            package_name = package_dir.name
            
            try:
                # 统计JS文件数量
                js_count, js_files = self.count_js_files(package_dir)
                
                # 记录包信息
                package_info = {
                    "package_name": package_name,
                    "js_count": js_count,
                    "js_files": [str(f.relative_to(package_dir)) for f in js_files],
                    "status": "processed"
                }
                
                if js_count == 1:
                    # 单个JS文件的包
                    if self.copy_package(package_dir, self.single_js_dir, "single"):
                        self.stats["single_js_packages"] += 1
                        package_info["category"] = "single_js"
                        print(f"  📄 {package_name}: 1个JS文件")
                    else:
                        self.stats["error_packages"] += 1
                        package_info["status"] = "copy_error"
                        
                elif js_count == 2:
                    # 两个JS文件的包
                    if self.copy_package(package_dir, self.double_js_dir, "double"):
                        self.stats["double_js_packages"] += 1
                        package_info["category"] = "double_js"
                        print(f"  📄📄 {package_name}: 2个JS文件")
                    else:
                        self.stats["error_packages"] += 1
                        package_info["status"] = "copy_error"
                        
                else:
                    # 跳过复杂包
                    self.stats["skipped_packages"] += 1
                    package_info["category"] = "skipped"
                    package_info["status"] = "skipped_complex"
                    print(f"  ⏭️ {package_name}: {js_count}个JS文件 (跳过)")
                
                self.extraction_log.append(package_info)
                
            except Exception as e:
                print(f"❌ 处理包失败 {package_name}: {e}")
                self.stats["error_packages"] += 1
                
                self.extraction_log.append({
                    "package_name": package_name,
                    "js_count": 0,
                    "js_files": [],
                    "status": "processing_error",
                    "error": str(e)
                })

    def save_extraction_report(self):
        """保存提取报告"""
        report = {
            "extraction_summary": {
                "source_directory": str(self.source_dir),
                "output_directory": str(self.output_dir),
                "extraction_time": self.get_current_time(),
                "statistics": self.stats
            },
            "package_details": self.extraction_log
        }
        
        report_file = self.output_dir / "extraction_report.json"
        
        try:
            with open(report_file, 'w', encoding='utf-8') as f:
                json.dump(report, f, indent=2, ensure_ascii=False)
            
            print(f"📋 提取报告已保存: {report_file}")
            
        except Exception as e:
            print(f"❌ 保存报告失败: {e}")

    def print_summary(self):
        """打印提取总结"""
        print("\n" + "="*60)
        print("📊 提取总结")
        print("="*60)
        print(f"总处理包数: {self.stats['total_packages']}")
        print(f"单JS文件包: {self.stats['single_js_packages']}")
        print(f"双JS文件包: {self.stats['double_js_packages']}")
        print(f"跳过复杂包: {self.stats['skipped_packages']}")
        print(f"处理错误包: {self.stats['error_packages']}")
        print("-" * 60)
        print(f"提取成功率: {((self.stats['single_js_packages'] + self.stats['double_js_packages']) / max(self.stats['total_packages'], 1) * 100):.1f}%")
        
        # 文件夹位置信息
        print(f"\n📁 输出目录结构:")
        print(f"  单JS包目录: {self.single_js_dir}")
        print(f"  双JS包目录: {self.double_js_dir}")
        
        # 统计各目录下的包数量
        single_count = len([d for d in self.single_js_dir.iterdir() if d.is_dir()]) if self.single_js_dir.exists() else 0
        double_count = len([d for d in self.double_js_dir.iterdir() if d.is_dir()]) if self.double_js_dir.exists() else 0
        
        print(f"  实际复制的单JS包: {single_count} 个")
        print(f"  实际复制的双JS包: {double_count} 个")

    def get_current_time(self):
        """获取当前时间字符串"""
        import datetime
        return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    def run(self):
        """运行完整的提取流程"""
        try:
            # 执行提取
            self.extract_simple_packages()
            
            # 保存报告
            self.save_extraction_report()
            
            # 打印总结
            self.print_summary()
            
            print(f"\n✅ 简单包提取完成！")
            
        except Exception as e:
            print(f"❌ 提取过程发生错误: {e}")


def main():
    """主函数"""
    # 创建提取器实例
    extractor = SimplePackageExtractor(
        source_dir="E:/extracted_packages/malicious",
        output_dir="E:/simplePackages"
    )
    
    # 运行提取
    extractor.run()


if __name__ == "__main__":
    main()